#ifndef INCLUDE_SHM_QUEUE_H
#define INCLUDE_SHM_QUEUE_H

#include <stdint.h>

#ifndef SHM_API
#ifdef SHM_STATIC
#define SHM_API static
#else
#ifdef __cplusplus
#define SHM_API extern "C"
#else
#define SHM_API extern
#endif
#endif
#endif

enum {
  SHM_QUEUE_ERROR_OK,
  SHM_QUEUE_ERROR_INVALID_PARAM = -1,
  SHM_QUEUE_ERROR_OPEN = -2,
  SHM_QUEUE_ERROR_NO_DATA = -3,
  SHM_QUEUE_ERROR_INVALID_DATA = -4,
  SHM_QUEUE_ERROR_NOT_SUPPORTED = -5,
  SHM_QUEUE_ERROR_TIMEOUT = -6
};

typedef enum shm_queue_mode {
  SHM_QUEUE_MODE_READ = 0x1,
  SHM_QUEUE_MODE_WRITE = 0x2,
  SHM_QUEUE_MODE_CREATE = 0x4,
  SHM_QUEUE_MODE_LOCK = 0x8,
} shm_queue_mode_t;

typedef struct shm_queue_options {
  const char *name;
  uint32_t mode;
  uint32_t elem_size;
  uint32_t elem_count;
} shm_queue_options_t;

typedef struct shm_queue {
  int fd;
  uint32_t mode;
  uint8_t *mem;
  uint32_t size;
  uint32_t elem_size;
  uint32_t elem_count;
  uint32_t rmsgid;
} shm_queue_t;

SHM_API int shm_queue_open(shm_queue_t *queue, const shm_queue_options_t *opt);
SHM_API int shm_queue_close(shm_queue_t *queue);

SHM_API int shm_queue_write(shm_queue_t *queue, void *data, uint32_t size);
SHM_API int shm_queue_read(shm_queue_t *queue, void *data, uint32_t size);

// read data from shared memory queue without copy data.
SHM_API int shm_queue_ncread(shm_queue_t *queue, void **data, uint32_t *size);

// write data to shared memory queue without copy data.
SHM_API int shm_queue_ncwrite_begin(shm_queue_t *queue, void **buf,
                                    uint32_t *buf_size);
SHM_API int shm_queue_ncwrite_finish(shm_queue_t *queue, uint32_t size);

#endif // INCLUDE_SHM_QUEUE_H

#ifdef SHM_QUEUE_IMPLEMENTATION
#include <errno.h>
#include <fcntl.h>
#include <pthread.h>
#include <stdbool.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <time.h>
#include <unistd.h>

#define ALIGN_TO_EIGHT __attribute__((aligned(8)))

typedef struct header {
  uint32_t size;
  uint32_t lastid;
  uint32_t lock;
  uint8_t padding[16];
} header_t ALIGN_TO_EIGHT;

typedef struct element {
  uint32_t size;
  uint32_t id;
  uint8_t data[0];
} element_t ALIGN_TO_EIGHT;

static inline element_t *_shm_queue_get_element(shm_queue_t *queue,
                                                uint32_t id) {
  return (element_t *)((uint8_t *)queue->mem + sizeof(header_t) +
                       (id % queue->elem_count) * queue->elem_size);
}

SHM_API int shm_queue_open(shm_queue_t *queue, const shm_queue_options_t *opt) {
  int fd;
  uint32_t size;
  struct stat stat;
  header_t* header;

  if (!queue || !opt || !opt->name || !opt->elem_count || !opt->elem_size) {
    return SHM_QUEUE_ERROR_INVALID_PARAM;
  }

  fd = -1;
  size = opt->elem_size * opt->elem_count + sizeof(header_t);

  if (opt->mode & SHM_QUEUE_MODE_WRITE || opt->mode & SHM_QUEUE_MODE_CREATE) {
    bool new_segment = false;
    while (true) {
      new_segment = true;
      fd = shm_open(opt->name, O_RDWR | O_CREAT | O_EXCL, 0644);
      if (fd >= 0) {
        fchmod(fd, 0644);
      } else if (errno == EEXIST) {
        fd = shm_open(opt->name, O_RDWR, 0644);
        if (fd < 0 && errno == ENOENT) {
          // the memory segment was deleted in the mean time
          usleep(1000);
          continue;
        }
        new_segment = false;
      } else {
        return SHM_QUEUE_ERROR_OPEN;
      }
      break;
    }

    if (new_segment) {
      // Allocate an extra `alignment` bytes as padding
      int result = ftruncate(fd, size);
      if (result == EINVAL) {
        return SHM_QUEUE_ERROR_OPEN;
      }
    }

    queue->elem_size = opt->elem_size;
    queue->elem_count = opt->elem_count;
    queue->mode = opt->mode;
    queue->fd = fd;
    queue->size = size;
    queue->rmsgid = 0;
    queue->mem =
        (uint8_t *)mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
    if (!queue->mem) {
      close(fd);
      return SHM_QUEUE_ERROR_OPEN;
    }
    if (new_segment) {
      header_t *header = (header_t *)queue->mem;
      header->lastid = 0;
      header->size = size;
      __atomic_store_n(&header->lock, 0, __ATOMIC_RELAXED);
    } else {
      header = (header_t*)queue->mem;
      queue->rmsgid = header->lastid;
    }
    return SHM_QUEUE_ERROR_OK;
  } else {
    fd = shm_open(opt->name, O_RDONLY, 0644);
    if (fd < 0) {
      return SHM_QUEUE_ERROR_OPEN;
    }

    if (fstat(fd, &stat) == 0) {
      size = stat.st_size;
    } else {
      return SHM_QUEUE_ERROR_OPEN;
    }

    queue->elem_size = opt->elem_size;
    queue->elem_count = opt->elem_count;
    queue->mode = opt->mode;
    queue->fd = fd;
    queue->size = size;
    queue->rmsgid = 0;
    queue->mem = (uint8_t *)mmap(NULL, size, PROT_READ, MAP_SHARED, fd, 0);
    if (!queue->mem) {
      close(fd);
      return SHM_QUEUE_ERROR_OPEN;
    }
    header = (header_t*)queue->mem;
    queue->rmsgid = header->lastid;

    return SHM_QUEUE_ERROR_OK;
  }
}

SHM_API int shm_queue_close(shm_queue_t *queue) {
  if (queue->mem) {
    munmap(queue->mem, queue->size);
    queue->mem = NULL;
  }
  if (queue->fd > 0) {
    close(queue->fd);
    queue->fd = -1;
  }
  return SHM_QUEUE_ERROR_OK;
}

SHM_API int shm_queue_test_set_id(shm_queue_t *queue, uint32_t new_id) {
  header_t *header;
  element_t *element;

  if (!queue) {
    return SHM_QUEUE_ERROR_INVALID_PARAM;
  }

  if (!queue->mem || (queue->mode & SHM_QUEUE_MODE_WRITE) == 0) {
    return SHM_QUEUE_ERROR_NOT_SUPPORTED;
  }

  header = (header_t *)queue->mem;
  header->lastid = new_id;
  return SHM_QUEUE_ERROR_OK;
}

SHM_API int shm_queue_lock(shm_queue_t *queue, long timeout_ms) {
  struct timespec start, now;
  header_t *header;
  if (!queue) {
    return SHM_QUEUE_ERROR_INVALID_PARAM;
  }
  header = (header_t *)queue->mem;
  clock_gettime(CLOCK_MONOTONIC, &start);

  long timeout_ns = timeout_ms * 1000000L;
  while (1) {
    int expected = 0;
    // 使用原子比较交换操作
    if (__atomic_compare_exchange_n(&header->lock, &expected, 1, 0,
                                    __ATOMIC_ACQUIRE, __ATOMIC_RELAXED)) {
      return 0; // 获取锁成功
    }

    clock_gettime(CLOCK_MONOTONIC, &now);
    long elapsed_ns = (now.tv_sec - start.tv_sec) * 1000000000L +
                      (now.tv_nsec - start.tv_nsec);

    if (elapsed_ns >= timeout_ns) {
      return SHM_QUEUE_ERROR_TIMEOUT;
    }

    // 硬件优化指令
#if defined(__x86_64__)
    __asm__ __volatile__("pause");
#elif defined(__aarch64__)
    __asm__ __volatile__("yield");
#endif
    sched_yield();
  }
}

SHM_API int shm_queue_unlock(shm_queue_t *queue) {
  header_t *header;
  if (!queue) {
    return SHM_QUEUE_ERROR_INVALID_PARAM;
  }
  header = (header_t *)queue->mem;
  __atomic_store_n(&header->lock, 0, __ATOMIC_RELEASE);
  return SHM_QUEUE_ERROR_OK;
}

SHM_API int shm_queue_write(shm_queue_t *queue, void *data, uint32_t size) {
  header_t *header;
  uint32_t new_id;
  element_t *element;

  if (!queue || !data || (size + sizeof(element_t) > queue->elem_size)) {
    return SHM_QUEUE_ERROR_INVALID_PARAM;
  }

  if (!queue->mem || (queue->mode & SHM_QUEUE_MODE_WRITE) == 0) {
    return SHM_QUEUE_ERROR_NOT_SUPPORTED;
  }

  header = (header_t *)queue->mem;
  if (queue->mode & SHM_QUEUE_MODE_LOCK) {
    shm_queue_lock(queue, 1000);
  }
  new_id = header->lastid + 1;
  element = _shm_queue_get_element(queue, new_id);
  element->id = new_id;
  element->size = size;
  memcpy(element->data, data, size);
  header->lastid = new_id;
  if (queue->mode & SHM_QUEUE_MODE_LOCK) {
    shm_queue_unlock(queue);
  }
  return SHM_QUEUE_ERROR_OK;
}

SHM_API int shm_queue_ncwrite_begin(shm_queue_t *queue, void **buf,
                                    uint32_t *buf_size) {
  header_t *header;
  uint32_t new_id;
  element_t *element;

  if (!queue) {
    return SHM_QUEUE_ERROR_INVALID_PARAM;
  }

  if (!queue->mem || (queue->mode & SHM_QUEUE_MODE_WRITE) == 0) {
    return SHM_QUEUE_ERROR_NOT_SUPPORTED;
  }

  if (queue->mode & SHM_QUEUE_MODE_LOCK) {
    shm_queue_lock(queue, 1000); 
  }

  header = (header_t *)queue->mem;
  new_id = header->lastid + 1;
  element = _shm_queue_get_element(queue, new_id);

  if (buf) {
    *buf = element->data;
  }

  if (buf_size) {
    *buf_size = queue->elem_size - sizeof(element_t);
  }
  return SHM_QUEUE_ERROR_OK;
}

SHM_API int shm_queue_ncwrite_finish(shm_queue_t *queue, uint32_t size) {
  header_t *header;
  uint32_t new_id;
  element_t *element;

  if (!queue || (size + sizeof(element_t) > queue->elem_size)) {
    return SHM_QUEUE_ERROR_INVALID_PARAM;
  }

  if (!queue->mem || (queue->mode & SHM_QUEUE_MODE_WRITE) == 0) {
    return SHM_QUEUE_ERROR_NOT_SUPPORTED;
  }

  header = (header_t *)queue->mem;
  new_id = header->lastid + 1;
  element = _shm_queue_get_element(queue, new_id);

  element->id = new_id;
  element->size = size;
  header->lastid = new_id;
  if (queue->mode & SHM_QUEUE_MODE_LOCK) {
    shm_queue_unlock(queue);
  }
  return SHM_QUEUE_ERROR_OK;
}

SHM_API int shm_queue_read(shm_queue_t *queue, void *data, uint32_t size) {
  header_t *header;
  uint32_t curid, nextid, lastid, osize;
  element_t *element;

  if (!queue || !data || !size) {
    return SHM_QUEUE_ERROR_INVALID_PARAM;
  }

  if (!queue->mem || (queue->mode & SHM_QUEUE_MODE_READ) == 0) {
    return SHM_QUEUE_ERROR_NOT_SUPPORTED;
  }

  header = (header_t *)queue->mem;
  curid = queue->rmsgid;
  lastid = header->lastid;

  if (lastid == curid) {
    return SHM_QUEUE_ERROR_NO_DATA;
  }

  nextid = curid + 1;
  element = _shm_queue_get_element(queue, nextid);

  // 检查元素ID是否匹配预期的nextid，解决溢出问题
  if (element->id != nextid) {
    nextid = lastid;// 调整到最新数据位置
    element = _shm_queue_get_element(queue, nextid);
  }

  if (element->size + sizeof(element_t) > queue->elem_size) {
    return SHM_QUEUE_ERROR_INVALID_DATA;
  }

  if (element->size > size) {
    return SHM_QUEUE_ERROR_INVALID_PARAM;
  }

  osize = element->size;
  memcpy(data, element->data, osize);
  queue->rmsgid = nextid; // 更新读取位置为当前元素的ID
  return osize;
}

SHM_API int shm_queue_ncread(shm_queue_t *queue, void **data, uint32_t *size) {
  header_t *header;
  uint32_t curid, nextid, lastid;
  element_t *element;

  if (!queue || !data || !size) {
    return SHM_QUEUE_ERROR_INVALID_PARAM;
  }

  if (!queue->mem || (queue->mode & SHM_QUEUE_MODE_READ) == 0) {
    return SHM_QUEUE_ERROR_NOT_SUPPORTED;
  }

  header = (header_t *)queue->mem;
  curid = queue->rmsgid;
  lastid = header->lastid;

  if (lastid == curid) {
    return SHM_QUEUE_ERROR_NO_DATA;
  }

  nextid = curid + 1;
  element = _shm_queue_get_element(queue, nextid);

  // 检查元素ID是否匹配预期的nextid
  if (element->id != nextid) {
    nextid = lastid;// 调整到最新数据位置
    element = _shm_queue_get_element(queue, nextid);
  }

  if (element->size + sizeof(element_t) > queue->elem_size) {
    return SHM_QUEUE_ERROR_INVALID_DATA;
  }

  *data = element->data;
  *size = element->size;
  queue->rmsgid = nextid; // 更新读取位置
  return SHM_QUEUE_ERROR_OK;
}
#undef SHM_QUEUE_IMPLEMENTATION
#endif // SHM_QUEUE_IMPLEMENTATION
